# -*- coding: utf-8 -*-
"""
Created on 2018 3.26
@author: hugh
"""
import numpy as np
import os
import config
import pickle
import dlib
import cv2
import data_aug


# 适用于二级目录 。。。/图片类别文件/图片（jpg）
class LoadImage(object):
	
	def __init__(self, img_dir, train_data_dir=config.processed_data_dir):
		"""初始化		
		Args:
			img_dir: 原始人脸图片目录的路径
			train_data_dir： 预处理后的图片路径
		"""
		if not os.path.exists(train_data_dir):
			self._detect_save_face(img_dir, train_data_dir)
		self.imgDir = train_data_dir

	def _detect_save_face(self, img_dir, train_data_dir):
		"""人脸检测、裁剪、对齐、保存处理后的图片		
		Args:
			img_dir: 原始人脸图片目录的路径
			train_data_dir： 预处理后的图片路径
		"""
		# 使用dlib自带的frontal_face_detector作为我们的特征提取器
		detector = dlib.get_frontal_face_detector()
		if config.alignment:
			sp = dlib.shape_predictor(config.predictor_path)
		index = 1
		size = min(config.IMAGE_WIDTH, config.IMAGE_HEIGHT)
		for (path, dirnames, filenames) in os.walk(img_dir):
			targetDir = train_data_dir + path[path.rfind('/'):]
			if len(filenames) != 0 and not os.path.exists(targetDir):
				os.makedirs(targetDir)
			for filename in filenames:
				if filename.endswith('.jpg'):
					print('Being processed picture %s' % index)
					img_path = path + '/' + filename
					# 从文件读取图片
					img = cv2.imread(img_path)
					# 转为rgb图片
					rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
					# 使用detector进行人脸检测 dets为返回的结果
					dets = detector(rgb_img, 1)
					num_faces = len(dets)
					# 去除干扰的图片（当一张图中检测到有多于一个人脸）
					if num_faces > 1 or num_faces == 0:
						continue
					if config.alignment:
						#人脸对齐并调整尺寸
						for d in dets:
							shape = sp(rgb_img, d)
						image = dlib.get_face_chip(rgb_img, shape, size)
						face = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
					else:
						#简单的切割和调整人脸尺寸
						# left：人脸左边距离图片左边界的距离 ；right：人脸右边距离图片左边界的距离
						# top：人脸上边距离图片上边界的距离 ；bottom：人脸下边距离图片上边界的距离
						for d in dets:
							x1 = d.top() if d.top() > 0 else 0
							y1 = d.bottom() if d.bottom() > 0 else 0
							x2 = d.left() if d.left() > 0 else 0
							y2 = d.right() if d.right() > 0 else 0
							# img[y:y+h,x:x+w]
							face = img[x1:y1, x2:y2]
							# 调整图片的尺寸
							face = cv2.resize(face, (config.IMAGE_WIDTH, config.IMAGE_HEIGHT))
					# 保存图片
					cv2.imwrite(targetDir + '/' + str(index) + '.jpg', face)
					index += 1

	def _load_img_path(self, foldName, img_label):
		"""加载图片路径和标签
		Args:
			foldName: 文件夹名或人脸类别名
			img_label： 图片对应的标签
		"""
		imgs = os.listdir(os.path.join(self.imgDir, foldName))
		imgNum = len(imgs)
		data_path = []
		labels = []
		for i in range (imgNum):
			img_path = os.path.join(self.imgDir, foldName, imgs[i])
			data_path.append(img_path)
			labels.append(int(img_label))
		return data_path, labels

	def _shuffle_train_data(self):
		"""数据集转为多维数组并打乱"""
		index = [i for i in range(len(self.train_imgs))]
		np.random.shuffle(index)
		self.train_imgs = np.asarray(self.train_imgs)
		self.train_labels = np.asarray(self.train_labels)
		self.train_imgs = self.train_imgs[index]
		self.train_labels = self.train_labels[index]

	def load_train_database_path(self):
		"""加载图片数据路径，并以文件夹名对应的索引号作为图片数据的标签，将映射字典保存"""
		img_path = os.listdir(self.imgDir)
		self.train_imgs = []
		self.train_labels = []
		self.dict = {}
		for i, foldName in enumerate(img_path):			
			data_path, labels = self._load_img_path(foldName, i)
			self.train_imgs.extend(data_path)
			self.train_labels.extend(labels)
			self.dict.setdefault(i, foldName)
			print ("文件名对应的label:")
			print (foldName, i)
		#保存训练标签和人名的字典
		save_train_label(self.dict)
		# 保存训练的类别数
		self.num_classes = len(self.dict)
		# 打乱数据集
		self._shuffle_train_data()
		# 数据集的数量
		self.image_n = len(self.train_imgs)
	
	def gen_train_valid_image(self, train_rate=config.train_rate):
		"""生成训练和验证数据
		Args:
			train_rate: 训练数据占比
		"""
		self.train_rate = train_rate
		self.load_train_database_path()
		self.train_n = int(self.image_n * self.train_rate)
		self.valid_n = int(self.image_n * (1 - self.train_rate))
		return self.train_imgs[0:self.train_n], self.train_labels[0:self.train_n], self.train_imgs[self.train_n:self.image_n], self.train_labels[self.train_n:self.image_n]

def save_train_label(train_label):
	"""保存训练标签和人名的对应关系"""
	with open(config.train_label, 'wb') as f:
		pickle.dump(train_label,f)

def read_train_label():
	"""读取标签和人名的字典"""
	with open(config.train_label, 'rb') as f:
		train_label = pickle.load(f)
	return train_label

def shuffle_train_data(train_imgs, train_labels):
	"""打乱数据"""
	index = [i for i in range(len(train_imgs))]
	np.random.shuffle(index)
	train_imgs = np.asarray(train_imgs)
	train_labels = np.asarray(train_labels)
	train_imgs = train_imgs[index]
	train_labels = train_labels[index]
	return train_imgs, train_labels

def get_next_batch_from_path(image_path, image_labels, pointer, IMAGE_HEIGHT=299, IMAGE_WIDTH=299, batch_size=64, is_train=True):
	"""批量生成图像"""
	batch_x = np.zeros([batch_size, IMAGE_HEIGHT,IMAGE_WIDTH,3])
	num_classes = len(image_labels[0])
	batch_y = np.zeros([batch_size, num_classes]) 
	for i in range(batch_size):	 
		image = cv2.imread(image_path[i+pointer*batch_size])
		image = cv2.resize(image, (IMAGE_HEIGHT, IMAGE_WIDTH))
		if is_train:
			img_aug = data_aug.data_aug(image)
			image = img_aug.get_aug_img()
		image = image / 255.0
		image = image - 0.5
		image = image * 2
		
		batch_x[i,:,:,:] = image
		batch_y[i] = image_labels[i+pointer*batch_size]
	return batch_x, batch_y

def gen_test_images(filename):
	"""从单个图片文件预处理后用于测试识别用"""
	size = min(config.IMAGE_WIDTH, config.IMAGE_HEIGHT)
	detector = dlib.get_frontal_face_detector()
	if config.alignment:
		sp = dlib.shape_predictor(config.predictor_path)
	if filename.endswith('.jpg'):
		print('Being processed picture %s' % filename)
		# 从文件读取图片
		img = cv2.imread(filename)
		# 转为rgb图片
		rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
		# 使用detector进行人脸检测 dets为返回的结果
		dets = detector(rgb_img, 1)		
		batch_x = np.zeros([len(dets), config.IMAGE_HEIGHT, config.IMAGE_WIDTH, 3])
		if config.alignment:
			i = 0
			#人脸对齐并调整尺寸			
			for d in dets:
				shape = sp(rgb_img, d)
				# Get the aligned face images
				# Optionally: 
				image_tmp = dlib.get_face_chip(rgb_img, shape, size)
				face = cv2.cvtColor(image_tmp, cv2.COLOR_RGB2BGR)
				image = face / 255.0
				image = image - 0.5
				image = image * 2
				batch_x[i, :, :, :] = image
				i += 1
		else:		
			#简单的切割和调整人脸尺寸
			# left：人脸左边距离图片左边界的距离 ；right：人脸右边距离图片左边界的距离
			# top：人脸上边距离图片上边界的距离 ；bottom：人脸下边距离图片上边界的距离
			for i, d in enumerate(dets):
				x1 = d.top() if d.top() > 0 else 0
				y1 = d.bottom() if d.bottom() > 0 else 0
				x2 = d.left() if d.left() > 0 else 0
				y2 = d.right() if d.right() > 0 else 0				
				face = img[x1:y1, x2:y2]
				# 调整图片的尺寸
				face = cv2.resize(face, (size, size))
				image = face / 255.0
				image = image - 0.5
				image = image * 2
				batch_x[i, :, :, :] = image
		return batch_x, dets